/te-ci pytorch L1

Greptile Summary

This PR fixes a subtle bug in get_reduction_params() within the debug utilities where passing tp_group=None (the case when TP is disabled, i.e. tp_size=1) would inadvertently assign None to reduction_group. Since group=None in PyTorch distributed collectives resolves to the default world process group, this caused unintended cross-rank reduction of weight tensor statistics across all ranks rather than restricting it to the intended TP group.

Key changes:

• In transformer_engine/debug/features/utils/__init__.py, the get_reduction_params() function now guards the reduction_group = tp_group assignment behind a tp_group is not None check (line 26).
• When weight_tensor_tp_group_reduce=True but tp_group is None, skip_reduction is set to True (line 30), preventing the accidental world-group reduction entirely.
• A clear inline comment (lines 24-25) explains why None must not be assigned to reduction_group.
• All callers (log_tensor_stats.py, log_fp8_tensor_stats.py, log_nvfp4_tensor_stats.py) remain unaffected as the API contract (skip_reduction, reduction_group, reduce_within_microbatch) is unchanged.

Confidence Score: 5/5

• This PR is safe to merge; it's a minimal, targeted bug fix that prevents real world-group reduction issues when TP is disabled.
• The fix is a single, well-scoped guard check that prevents a concrete bug (world-group reduction when tp_group=None) without altering any existing code paths. The logic is straightforward, inline comments clearly explain the intent, and the change is localized to a utility function with a stable API contract.
• No files require special attention.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["get_reduction_params(tensor_name, tp_group)"] --> B{tensor_name == 'weight'?}
    B -- No --> C["skip_reduction=False<br/>reduction_group=default<br/>reduce_within_microbatch=True"]
    B -- Yes --> D{weight_tensor_tp_group_reduce?}
    D -- False --> E["skip_reduction=True<br/>reduce_within_microbatch=False"]
    D -- True --> F{tp_group is not None?}
    F -- Yes --> G["reduction_group=tp_group<br/>skip_reduction=False<br/>reduce_within_microbatch=False"]
    F -- No --> H["skip_reduction=True ✅ FIXED<br/>reduce_within_microbatch=False<br/>(was: reduction_group=None → world group 🐛)"]

_{Last reviewed commit: c2adcaf}